package main

import (
	"fmt"
	"gitee.com/simple-set/simple.io/src/event"
	"gitee.com/simple-set/simple.io/src/process/codec"
	"gitee.com/simple-set/simple.io/src/process/socks"
	"gitee.com/simple-set/simple.io/src/socket"
	"io"
	"log/slog"
	"net"
	"reflect"
	"strings"
)

type AuthUserPass struct {
}

func NewAuthUserPass() *AuthUserPass {
	return &AuthUserPass{}
}

func (a AuthUserPass) AuthMethod() socks.AuthType {
	return socks.UserPassAuth
}

func (a AuthUserPass) Challenge(user, pass string) bool {
	if user == "admin" && pass == "123456" {
		return true
	}
	return false
}

type baiduInterceptor struct{}

func newBaiduInterceptor() *baiduInterceptor {
	return &baiduInterceptor{}
}

func (b baiduInterceptor) Intercept(request *socks.Request) bool {
	if request.DstAddrType == socks.Domain {
		if strings.HasPrefix(string(request.DstAddr), "www.baidu.com") {
			return true
		}
	}
	return false
}

type ConnectHandler struct {
	client *event.BootClient
}

func NewConnectHandler() *ConnectHandler {
	bootClient := event.NewBootClient()
	bootClient.Register(newSocksBackend())
	return &ConnectHandler{client: bootClient}
}

func (c *ConnectHandler) Connect(request *socks.Request, response *socks.Response) error {
	starterClient, err := c.getConnect(request)
	if err != nil {
		_ = response.Reply(socks.HostUnreachable)
		return err
	}
	defer func() {
		if starterClient != nil {
			_ = starterClient.Close()
		}
	}()

	if err = response.SetAddrAndPort(starterClient.Session().Socket.LocalAddr()); err != nil {
		_ = response.Reply(socks.ServerFailure)
		return err
	}
	if err = response.Reply(socks.SuccessReply); err != nil {
		return err
	}
	if _, err = starterClient.Request(request); err != nil {
		return err
	}
	return io.EOF
}

func (c *ConnectHandler) getConnect(request *socks.Request) (event.StarterClient, error) {
	var serverAddr string
	if request.DstAddrType == socks.Domain {
		serverAddr = fmt.Sprintf("%s:%d", string(request.DstAddr), request.DstPort)
	} else {
		serverAddr = request.GetDstAddr().String()
	}

	if starterClient, err := c.client.DialTcp(serverAddr); err != nil {
		return nil, err
	} else {
		return starterClient, nil
	}
}

type socksBackend struct {
}

func newSocksBackend() *socksBackend {
	return &socksBackend{}
}

func (s socksBackend) Active(session *event.Session, value any) (any, error) {
	request := value.(*socks.Request)
	go codec.NewReadByteBuf(session.Socket.Reader).WriteTo(request.Buff())
	return request.Buff().WriteTo(session.Socket.Writer)
}

type AssociateHandler struct {
	bootClient *event.BootClient
}

func NewAssociateHandler() *AssociateHandler {
	return &AssociateHandler{bootClient: event.NewBootClient()}
}

func (a *AssociateHandler) Disconnect(session *event.Session) error {
	if udpServer := session.GetAttr("udpServer"); udpServer != nil {
		if server, ok := udpServer.(socket.Server); ok {
			_ = server.Close()
		}
	}
	return nil
}

func (a *AssociateHandler) Associate(request *socks.Request, response *socks.Response) error {
	if udpServer, err := a.createUdpServer(request, response); err != nil {
		return err
	} else {
		request.Session().SetAttr("udpServer", udpServer)
		go a.relay(udpServer, request)
	}
	return nil
}

func (a *AssociateHandler) relay(udpServer socket.Server, request *socks.Request) {
	for {
		sock, err := udpServer.Listen()
		if err != nil {
			_ = udpServer.Close()
			slog.Error(reflect.ValueOf(udpServer).Elem().Type().String(), err)
			return
		}

		udpRelayRequest := socks.NewUdpRelay(codec.NewReadByteBuf(sock))
		if err := udpRelayRequest.Decode(); err != nil {
			slog.Error("udpRelayRequest data decode failed", err)
			continue
		}
		if !a.udpRelayIsAllow(sock.RemoteAddr().(*net.UDPAddr), request) {
			slog.Warn("udpRelay Address mismatch", "dstAddr", udpRelayRequest.DstAddr)
			continue
		}
		udpRelayPayload, err := udpRelayRequest.Payload()
		if err != nil {
			slog.Error("udpRelayRequest.Payload failed", err)
			continue
		}

		// 转发udp中继的有效载荷到目标服务器
		receivedData, err := a.udpRelaTransfer(udpRelayRequest.GetDstAddr(), udpRelayPayload)
		if err != nil {
			slog.Error("udpRelaTransfer failed", err)
			continue
		} else {
			udpRelayResponse := socks.NewUdpRelay(codec.NewWriteByteBuf(sock))
			udpRelayResponse.SetPayload(receivedData)
			udpRelayResponse.SetAddr(udpRelayRequest.GetDstAddr())
			if err := udpRelayResponse.Encode(); err != nil {
				slog.Error("udpRelayRequest encode failed", err)
			}
		}
	}
}

// 由Associate指令触发，启动新UDP监听
func (a *AssociateHandler) createUdpServer(request *socks.Request, response *socks.Response) (socket.Server, error) {
	udpServer := socket.NewUdpServer(fmt.Sprintf("%s:", a.getListenAddr(request)))
	if err := udpServer.Bind(); err != nil {
		_ = response.Reply(socks.NetworkUnreachable)
		return nil, err
	}

	if err := response.SetAddrAndPort(udpServer.ListenAddr()); err != nil {
		_ = response.Reply(socks.ServerFailure)
		return nil, err
	}
	if err := response.Reply(socks.SuccessReply); err != nil {
		return nil, err
	}
	return udpServer, nil
}

func (a *AssociateHandler) createUdpClient(addr *net.UDPAddr) (event.StarterClient, error) {
	if client, err := a.bootClient.DialUdp(addr.String()); err != nil {
		return nil, err
	} else {
		return client, nil
	}
}

func (a *AssociateHandler) getListenAddr(request *socks.Request) string {
	if tcpAddr, ok := request.Session().Socket.LocalAddr().(*net.TCPAddr); ok {
		return tcpAddr.IP.String()
	}
	return request.Session().Socket.LocalAddr().(*net.UDPAddr).IP.String()
}

func (a *AssociateHandler) udpRelayIsAllow(relayAddr *net.UDPAddr, request *socks.Request) bool {
	if net.IP(request.DstAddr).IsUnspecified() {
		return relayAddr.IP.Equal(request.Session().Socket.RemoteAddr().(*net.TCPAddr).IP)
	} else {
		return relayAddr.IP.Equal(request.DstAddr)
	}
}

func (a *AssociateHandler) udpRelaTransfer(addr *net.UDPAddr, sendData []byte) ([]byte, error) {
	sock, err := socket.NewUdpClient().Dial(addr.String())
	if err != nil {
		return nil, err
	}
	if err := sock.WriterBytes(sendData); err != nil {
		return nil, err
	}

	// 从目标服务器读取回复, 受限MTU大小, UDP接收缓冲区1500字节
	receivedData := make([]byte, 1500)
	if n, err := sock.Read(receivedData); err != nil {
		return nil, err
	} else {
		return receivedData[:n], nil
	}
}
